package semantic_compute.common.entitylinking;

import mpi.aida.Disambiguator;
import mpi.aida.access.DataAccess;
import mpi.aida.config.settings.DisambiguationSettings;
import mpi.aida.config.settings.PreparationSettings;
import mpi.aida.config.settings.disambiguation.CocktailPartyDisambiguationSettings;
import mpi.aida.config.settings.preparation.StanfordHybridPreparationSettings;
import mpi.aida.data.*;
import mpi.aida.preparator.Preparator;

import java.util.HashSet;
import java.util.Map;
import java.util.Set;

/**
 * Created by julianzliu on 4/10/2017.
 */
public class AIDAEntityLinker {


    /**********************************************************************
     * UNIT TEST
     */
    public static void test_hybrid_mention(){
        // Define the input.
        String inputText = "When [[Page]] played Kashmir at Knebworth, his Les Paul was uniquely tuned.";
        try {
            // Prepare the input for disambiguation. The Stanford NER will be run
            // to identify names. Strings marked with [[ ]] will also be treated as names.
            PreparationSettings prepSettings = new StanfordHybridPreparationSettings();
            Preparator p = new Preparator();
            PreparedInput input = p.prepare(inputText, prepSettings);

            // Disambiguate the input with the graph coherence algorithm.
            DisambiguationSettings disSettings = new CocktailPartyDisambiguationSettings();
            Disambiguator d = new Disambiguator(input, disSettings);
            DisambiguationResults results = d.disambiguate();

            // Print the disambiguation results.
            for (ResultMention rm : results.getResultMentions()) {
                ResultEntity re = results.getBestEntity(rm);
                System.out.println(rm.getMention() + " -> " + re);
            }

            Set<KBIdentifiedEntity> entities = new HashSet<KBIdentifiedEntity>();
            for (ResultMention rm : results.getResultMentions()) {
                entities.add(results.getBestEntity(rm).getKbEntity());
            }

            Map<KBIdentifiedEntity, EntityMetaData> entitiesMetaData =
                    DataAccess.getEntitiesMetaData(entities);

            for (ResultMention rm : results.getResultMentions()) {
                KBIdentifiedEntity entity = results.getBestEntity(rm).getKbEntity();
                EntityMetaData entityMetaData = entitiesMetaData.get(entity);

                if (Entities.isOokbEntity(entity)) {
                    System.out.println("\t" + rm + "\t NO MATCHING ENTITY");
                } else {
                    System.out.println("\t" + rm + "\t" + entityMetaData.getId() + "\t"
                            + entity + "\t" + entityMetaData.getHumanReadableRepresentation()
                            + "\t" + entityMetaData.getUrl());
                }
            }

        }catch (Exception e){
            e.printStackTrace();
        }
    }


    public static void main(String[] args){
        test_hybrid_mention();
    }

}
